#!/usr/bin/python -tt

import numpy as np
import pyfits

import time

import logging
import logging.config
import os
scriptpath = os.path.dirname(os.path.realpath(__file__))
logging.config.fileConfig(scriptpath+'/logging.conf')
logger = logging.getLogger('fitswriter')

########### DEPRICATED .. just here to show how it doesn't work ##############
def version1():
    
    EventNum = np.zeros(1, dtype=np.uint32)
    TriggerNum = np.zeros(1, dtype=np.uint32)
    TriggerType = np.zeros(1, dtype=np.uint32)
    NumBoards = np.zeros(1, dtype=np.uint32)
    Errors = np.zeros(4, dtype=np.uint8)
    SoftTrig = np.zeros(1, dtype=np.uint32)
    UnixTimeUTC = np.zeros(2, dtype=np.uint32)
    BoardTime = np.zeros(40, dtype=np.uint32)
    StartCellData = np.zeros(1440, dtype=np.uint16)
    StartCellTimeMarker = np.zeros(160, dtype=np.uint16)
    Data = np.zeros(1474560, dtype=np.int16)
    
    col1 = pyfits.Column(name='EventNum', format='1J', unit="uint32", array=EventNum)
    col2 = pyfits.Column(name='TriggerNum', format='1J', unit="uint32", array=TriggerNum )
    col3 = pyfits.Column(name='TriggerType', format='1I', unit="uint32", array=TriggerType )
    col4 = pyfits.Column(name='NumBoards', format='1J', unit="uint32", array=NumBoards )
    col5 = pyfits.Column(name='Errors', format='4B', unit="uint8", array=Errors )
    col6 = pyfits.Column(name='SoftTrig', format='1J', unit="uint32", array=SoftTrig )
    col7 = pyfits.Column(name='UnixTimeUTC', format='2J',unit="uint32", array=UnixTimeUTC )
    col8 = pyfits.Column(name='BoardTime', format='40J',unit="uint32", array=BoardTime )
    col9 = pyfits.Column(name='StartCellData', format='1440I',unit="uint16", array=StartCellData )
    col10 = pyfits.Column(name='StartCellTimeMarker', format='160I',unit="uint16", array=StartCellTimeMarker )
    col11 = pyfits.Column(name='Data', format='1474560I',unit="int16", array=Data )
    
    cols = pyfits.ColDefs([col1, col2, col3, col4, col5, col6, col7, col8, col9, col10,col11])
    
    tbhdu = pyfits.new_table(cols)
    hdu = pyfits.PrimaryHDU()
    thdulist = pyfits.HDUList([hdu, tbhdu])

    thdulist.writeto('table.fits')

########### DEPRICATED .. just here to show how it doesn't work ##############
def version2():
    col1 = pyfits.Column(name='EventNum', format='1J', unit="uint32")
    col2 = pyfits.Column(name='TriggerNum', format='1J', unit="uint32")
    col3 = pyfits.Column(name='TriggerType', format='1I', unit="uint32")
    col4 = pyfits.Column(name='NumBoards', format='1J', unit="uint32")
    col5 = pyfits.Column(name='Errors', format='4B', unit="uint8")
    col6 = pyfits.Column(name='SoftTrig', format='1J', unit="uint32")
    col7 = pyfits.Column(name='UnixTimeUTC', format='2J',unit="uint32")
    col8 = pyfits.Column(name='BoardTime', format='40J',unit="uint32")
    col9 = pyfits.Column(name='StartCellData', format='1440I',unit="uint16")
    col10 = pyfits.Column(name='StartCellTimeMarker', format='160I',unit="uint16")
    col11 = pyfits.Column(name='Data', format='1474560I',unit="int16")

    cols = pyfits.ColDefs([col1, col2, col3, col4, col5, col6, col7, col8, col9, col10, col11])

    mydata = np.recarray(1, dtype = np.dtype([
            ('EventNum', '<u4'), 
            ('TriggerNum', '<u4'), 
            ('TriggerType', '<u4'), 
            ('NumBoards', '<u4'), 
            ('Errors', '|u1', 4), 
            ('SoftTrig', '<u4'), 
            ('UnixTimeUTC', '<u4', 2), 
            ('BoardTime', '<u4', 40), 
            ('StartCellData', '<u2', 1440), 
            ('StartCellTimeMarker', '<u2', 160), 
            ('Data', '<i2', (1440,1024))]) )

    tbhdu = pyfits.new_table(cols)
    hdu = pyfits.PrimaryHDU()
    thdulist = pyfits.HDUList([hdu, tbhdu])

    tbhdu.data = mydata

    thdulist.writeto('table.fits')

########### DEPRICATED .. just here to show how it doesn't work ##############
def version3():
    
    EventNum = np.ones(1, dtype=np.uint32) * 0x11111111
    TriggerNum = np.ones(1, dtype=np.uint32) * 0x22222222
    TriggerType = np.ones(1, dtype=np.uint32) * 0x33333333
    NumBoards = np.ones(1, dtype=np.uint32) * 0x44444444
    Errors = np.ones(4, dtype=np.uint8) * 0x55
    SoftTrig = np.ones(1, dtype=np.uint32) * 0x66666666
    UnixTimeUTC = np.ones(2, dtype=np.uint32) * 0x77777777
    BoardTime = np.ones(4, dtype=np.uint32) * 0x88888888
    StartCellData = np.ones(14, dtype=np.uint16) * 0x9999
    StartCellTimeMarker = np.ones(16, dtype=np.uint16) * 0xaaaa
    #Data = np.zeros(1474560, dtype=np.int16)
    
    col1 = pyfits.Column(name='EventNum',       format='1J', unit="uint32", array=EventNum)
    col2 = pyfits.Column(name='TriggerNum',     format='1J', unit="uint32", array=TriggerNum )
    col3 = pyfits.Column(name='TriggerType',    format='1J', unit="uint32", array=TriggerType )
    col4 = pyfits.Column(name='NumBoards',      format='1J', unit="uint32", array=NumBoards )
    col5 = pyfits.Column(name='Errors',         format='4B', unit="uint8",  array=Errors )
    col6 = pyfits.Column(name='SoftTrig',       format='1J', unit="uint32", array=SoftTrig )
    col7 = pyfits.Column(name='UnixTimeUTC',    format='2J', unit="uint32", array=UnixTimeUTC )
    col8 = pyfits.Column(name='BoardTime',      format='4J', unit="uint32", array=BoardTime )
    col9 = pyfits.Column(name='StartCellData',  format='14I',unit="uint16", array=StartCellData )
    col10 = pyfits.Column(name='StartCellTimeMarker', format='16I',unit="uint16", array=StartCellTimeMarker )
    #col11 = pyfits.Column(name='Data', format='1474560I',unit="int16", array=Data )
    
    cols = pyfits.ColDefs([col1, col2, col3, col4, col5, col6, col7, col8, col9, col10])
    
    tbhdu = pyfits.new_table(cols)
    hdu = pyfits.PrimaryHDU()
    thdulist = pyfits.HDUList([hdu, tbhdu])

    thdulist.writeto('table.fits')


class RawDataFitsWriter( object ):
    
    headerTemplateFileName = "fitswriter_rawdata_header_lines.txt"
    
    def __init__( self, outfilepath, thedict ):
        logger.debug("RawDataFitsWriter __init__ called")
        self.f = open(outfilepath, 'wb')
        
        self.header_file = open(self.headerTemplateFileName)
        for line in self.header_file:
            # strip off \n at the end of each line
            line = line[:-1]
            # make line 80 characters long
            if len(line) > 80:
                line = line[:80]
            line += (80-len(line))*' '
            self.f.write(line)
        self.header_file.close()
        
        roi = thedict['settings']['SimInfo']['NumberOfSamples']
        self.rawdata = np.recarray(1, dtype = np.dtype([
            ('EventNum', '>u4'), 
            ('TriggerNum', '>u4'), 
            ('TriggerType', '>u2'), 
            ('NumBoards', '>u4'), 
            ('Errors', '|u1', 4), 
            ('SoftTrig', '>u4'), 
            ('UnixTimeUTC', '>u4', 2), 
            ('BoardTime', '>u4', 40), 
            ('StartCellData', '>u2', 1440), 
            ('StartCellTimeMarker', '>u2', 160), 
            ('Data', '>i2', (1440,roi))]) )

    def write( self ):
        logger.debug("RawDataFitsWriter writeRawData() called")
        self.f.write(self.rawdata)

    def __del__( self ):
        logger.debug("RawDataFitsWriter __del__ called")
        self.f.close()

class DrsCalibFitsWriter( object ):
    
    headerTemplateFileName = "fitswriter_drsfile_header_lines.txt"
    
    def __init__( self, outfilepath ):
        logger.debug("DrsCalibFitsWriter __init__ called")
        self.f = open(outfilepath, 'wb')
        
        self.header_file = open(self.headerTemplateFileName)
        for line in self.header_file:
            # strip off \n at the end of each line
            line = line[:-1]
            # make line 80 characters long
            if len(line) > 80:
                line = line[:80]
            line += (80-len(line))*' '
            self.f.write(line)
        self.header_file.close()
        
        self.drscalib = np.recarray(1, dtype = np.dtype([
            ('RunNumberBaseline', '<u4'), 
            ('RunNumberGain', '<u4'), 
            ('RunNumberTriggerOffset', '<u4'),
            ('BaselineMean', '<f4', (1440,1024)),
            ('BaselineRms', '<f4', (1440,1024)),
            ('GainMean', '<f4', (1440,1024)),
            ('GainRms', '<f4', (1440,1024)),
            ('TriggerOffsetMean', '<f4', (1440,1024)),
            ('TriggerOffsetRms', '<f4', (1440,1024)) ]) )
            # TriggerOffsetTMMean and
            # TriggerOffsetTMRms are left out, since they are empty.

    def write( self ):
        logger.debug("DrsCalibFitsWriter writeRawData() called")
        self.f.write(self.drscalib)

    def __del__( self ):
        logger.debug("DrsCalibFitsWriter __del__ called")
        self.f.close()

class RawDataFitsWriter2( object ):
    
    def __init__( self, outfilepath, thedict ):
        logger.debug("RawDataFitsWriter __init__ called")
        self.f = open(outfilepath, 'wb')
        
        self.creationtime = time.time()
        
        self.roi = thedict['settings']['SimInfo']['NumberOfSamples']
        self.run_number = thedict['settings']['run_number']
        self.run_type_string = thedict['settings']['run_type_string']
        roi = self.roi
        self.rawdata = np.recarray(1, dtype = np.dtype([
            ('EventNum', '>u4'), 
            ('TriggerNum', '>u4'), 
            ('TriggerType', '>u2'), 
            ('NumBoards', '>u4'), 
            ('Errors', '|u1', 4), 
            ('SoftTrig', '>u4'), 
            ('UnixTimeUTC', '>u4', 2), 
            ('BoardTime', '>u4', 40), 
            ('StartCellData', '>u2', 1440), 
            ('StartCellTimeMarker', '>u2', 160), 
            ('Data', '>i2', (1440,roi))]) )
            
        self.tformcodes = {
            'EventNum'     : '1J',
            'TriggerNum'   : '1J',
            'TriggerType'  : '1I',
            'NumBoards'    : '1J',
            'Errors'       : '4B',
            'SoftTrig'     : '1J',
            'UnixTimeUTC'  : '2J',
            'BoardTime'    : '40J',
            'StartCellData':'1440I',
            'StartCellTimeMarker': '160I',
            'Data'                  : str(roi*1440)+'I'
        }
        
        self.writeFirstHeader()
        self.starttime = time.time()
        self.stoptime = time.time()
        self.number_of_events = 0
        self.writeSecondHeader( self.starttime, self.stoptime, self.number_of_events )

    def write( self ):
        logger.debug("RawDataFitsWriter writeRawData() called")
        self.f.write(self.rawdata)
        self.number_of_events += 1
        self.stoptime = time.time()

    def __del__( self ):
        logger.debug("RawDataFitsWriter __del__ called")
        self.f.seek(0)
        self.writeFirstHeader()
        self.writeSecondHeader( self.starttime, self.stoptime, self.number_of_events )
        self.f.close()


    def writeFirstHeader( self ):
        
        firstheader = [
            ['simple', True, 'file does conform to FITS standard'],
            ['bitpix' , 8, 'number of bits per data pixel'],
            ['naxis' , 0, 'number of data axes'],
            ['extend' ,True , 'FITS dataset may contain extensions'],
            ['end'] ]
            
        self._writeSomeHeader( firstheader )
        
    def writeSecondHeader( self, starttime, stoptime, number_of_events):
        
        datatype = self.rawdata.dtype
        ctime = self.creationtime
        gmtime = time.gmtime(ctime)
        
        raw_data_width = datatype.itemsize
        number_of_fields = len(datatype.fields)
        #checksum = 'xxxx'
        #datasum = 'xxxxx'
        starttime = starttime/60/60/24
        stoptime = stoptime/60/60/24
        run_number = self.run_number
        run_type_string = self.run_type_string
        
        secondheader = [
        ['xtension', 'BINTABLE','binary table extension'],
        ['bitpix', 8, '8-bit bytes'],
        ['naxis', 2, '2-dimensional binary table' ],
        ['naxis1', raw_data_width, 'width of table in bytes'],
        ['naxis2', number_of_events, 'number of rows in table'],
        ['pcount', 0 ,'size of special data area'],
        ['gcount', 1, 'one data group (required keyword)'],
        ['tfields', number_of_fields, 'number of fields in each row'],
        ['extname', 'Events', 'name of extension table'],
        #['checksum', checksum, 'Checksum for the whole HDU'],
        #['datasum', datasum, 'Checksum for the data block'],
        ]
        for i,name in enumerate(datatype.names):
            secondheader.append(['ttype'+str(i), name])
            if datatype[name].subdtype is None:
                secondheader.append(['tunit'+str(i), datatype[name].name])
            else:
                secondheader.append(['tunit'+str(i), datatype[name].subdtype[0].name])
            secondheader.append(['tform'+str(i), self.tformcodes[name]])
            
        secondheader += [
        ['telescop', 'FACT', 'Telescope that acquired this data'],
        ['creator' , 'jafs-RawDataFitsWriter2' ],
        ['origin', 'FACT - TU Dortmund', 'Institution that wrote the file'],
        ['date', time.strftime("%Y-%m-%dT%H:%M:%S", gmtime), 'File creation date'],
        ['night', gmtime.tm_year*10000+gmtime.tm_mon*100+gmtime.tm_mday, 'Night as int'],
        ['timesys', 'UTC'],
        ['timeunit', 'd', 'Time given in days w.r.t. to MJDREF'],
        ['mjdref', 40587, 'MJD to UNIX time (seconds since 1970/1/1)'],
        ['runid', run_number, 'Run number'],
        ['nboard', 40, 'Number of acquisition boards'],
        ['npix', 1440, 'Number of pixels'],
        ['ntmark', 160, 'Number of time marker channels'],
        ['ncells', 1024, 'Maximum number of slices per pixels'],
        ['nroi', self.roi, 'Number of slices per pixels'],
        ['nroitm', 0, 'Number of slices per time-marker'],
        ['tmshift', 0, 'Shift of marker readout w.r.t. to data'],
        ['daq', 'simulated perfect ADC', 'Data acquisition type'],
        ['adcrange', 2000, 'Dynamic range in mV'],
        ['adc', 12, 'Resolution in bits'],
        ['runtype', run_type_string, 'File type according to FAD configuration'],
        ['refclk', 2048., 'Average reference clock frequency in Hz' ],
        ['drscalib', False, 'This file belongs to a DRS calibration'],
        ['tstarti', int(starttime), 'Time when first evt received (integral part)'],
        ['tstartf', starttime - int(starttime), 'Time when first evt received (integral part)'],
        ['tstopi', int(stoptime), 'Time when last evt received (integral part)'],
        ['tstopf', stoptime - int(stoptime), 'Time when last evt received (integral part)'],
        ['ntrg' , number_of_events, 'No of physics triggers written'],
        ['ntrgped', 0, 'No of pure pedestal triggers written'],
        ['ntrglpe', 0, 'No of external light pulser triggers written'],
        ['ntrgtim', 0, 'No of time calibration triggers written'],
        ['ntrglpi', 0, 'No of internal light pulser triggers written'],
        ['ntrgext1', 0,'No of triggers from ext1 written' ],
        ['ntrgext2', 0,'No of triggers from ext2 written' ],
        ['ntrgmisc', 0, 'No of all other triggers' ],
        ['end'] ]

        self._writeSomeHeader( secondheader )
        
    def _writeSomeHeader( self, listOfKeyValueComments ):
        header = ''
        
        for keyValueComment in listOfKeyValueComments:
            key, value, comment = None,None,None
            if len(keyValueComment) == 3:
                key, value, comment = keyValueComment
            elif len(keyValueComment) == 2:
                key, value = keyValueComment
            elif len(keyValueComment) == 1:
                key = keyValueComment[0]
            else:
                pass
                
            line = ''
            line += key.upper()[:8].ljust(8)
            
            # Add Value if present
            if value is not None:
                line += "= "
                # Check type of Value and add the string
                if type(value) == type(True):
                    if value:
                        line += 'T'
                    else:
                        line += 'F'
                elif type(value) == type(1):
                    line += str(value)[:20].rjust(20)
                elif type(value) == type(1.):
                    line += str(round(value,20))[:20].upper().rjust(20)
                elif type(value) == type(""):
                    maxlen = 80-len(line)-2
                    line += "'" + value[:maxlen].ljust(20) + "'"
                else:
                    raise Exception("unsupported type", type(value))
                    

            if comment is not None:
                if len(line)+3 < 80:
                    line += ' / '
                    maxlen = 80-len(line)
                    if type(comment) is type(''):
                        line += comment[:maxlen]
                    else:
                        line += str(comment)[:maxlen]

            line = line.ljust(80)
            if len(line) > 80:
                raise Exception("len line > 80")
            
            header += line
            
        header = header.ljust( (len(header)/2880+1)*2880)
        
        self.f.write(header)

if __name__ == '__main__':
    import rlcompleter
    import readline
    readline.parse_and_bind('tab: complete')

    fitsfile = RawDataFitsWriter("test.fits")
    #~ ('EventNum', '>u4'), 
    #~ ('TriggerNum', '<u4'), 
    #~ ('TriggerType', '<u4'), 
    #~ ('NumBoards', '<u4'), 
    #~ ('Errors', '|u1', 4), 
    #~ ('SoftTrig', '<u4'), 
    #~ ('UnixTimeUTC', '<u4', 2), 
    #~ ('BoardTime', '<u4', 40), 
    #~ ('StartCellData', '<u2', 1440), 
    #~ ('StartCellTimeMarker', '<u2', 160), 
    #~ ('Data', '>i2', (1440,1024))]) )
    
    fitsfile.rawdata['EventNum'] = 1
    fitsfile.rawdata['TriggerNum'] = 2
    fitsfile.rawdata['TriggerType'] = 3
    fitsfile.rawdata['NumBoards'] = 4
    fitsfile.rawdata['Errors'] = np.ones(4,dtype=np.uint8)*5
    fitsfile.rawdata['SoftTrig'] = 6
    fitsfile.rawdata['UnixTimeUTC'] = np.ones(2,dtype=np.uint32)*7
    fitsfile.rawdata['BoardTime'] = np.ones(40,dtype=np.uint32)*8
    fitsfile.rawdata['StartCellData'] = np.ones(1440,dtype=np.uint16)*9
    fitsfile.rawdata['StartCellTimeMarker'] = np.ones(160,dtype=np.uint16)*10
    fitsfile.rawdata['Data'] = np.ones((1440,1024),dtype=np.int16)*11
    
    fitsfile.write()
    del fitsfile
    
    print "wrote: test.fits"
    
    drsfile = DrsCalibFitsWriter("test.drs.fits")
    drsfile.drscalib['GainMean'] *= 4
    drsfile.write()
    del drsfile
    print "wrote: test.drs.fits"
